// Copyright (c) 2010, Google Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// ---
// Author: Jeffrey Junfeng Pan (panjunfeng@gmail.com)

/*

This is the core data structure an algorithm of frequent pattern itemset mining
using the frequent pattern tree algorithm.

For more detail of the algorithm, please refer to the following paper:

@article{Han2004FPTree,
  author = {Han, Jiawei and Pei, Jian and Yin, Yiwen and Mao, Runying},
  title = {Mining Frequent Patterns without Candidate Generation:
           A Frequent-Pattern Tree Approach},
  journal = {Data Min. Knowl. Discov.},
  volume = {8},
  number = {1},
  year = {2004},
  issn = {1384-5810},
  pages = {53--87},
  doi = {http://dx.doi.org/10.1023/B:DAMI.0000005258.31418.83},
  publisher = {Kluwer Academic Publishers},
  address = {Hingham, MA, USA},
}
*/

#ifndef FP_GROWTH_TREE_H_
#define FP_GROWTH_TREE_H_

#include <cmath>
#include <algorithm>
#include <unordered_map>
#include <string>
#include <vector>

#include "common/logging.h"
#include "common/macros.h"
#include "common/scoped_ptr.h"
#include "common/strutil.h"

#include "dictionary.h"
#include "header.h"
#include "itemsets.h"
#include "node.h"
#include "transactions.h"
#include "transactions_file_impl.h"
#include "transactions_memory_impl.h"
#include "pattern_writer.h"

namespace frequent_pattern_mining {

// ------------------------------------------------------------
//                      Declaration
// ------------------------------------------------------------

template <typename T>
class Tree {
 public:

  // Constructor
  explicit Tree(T impossible_item);

  // Destructor
  ~Tree();

  // Build the Frequent Pattern Tree using either min_support_count or
  // min_support_ratio. Only one min_support shall be properly set.
  // The other is set to any negative value. For example:
  // 1) Set by min_support_count = 10
  //    BuildTree(transactions, 10, -1);
  // 2) Set by min_support_ratio = 0.5
  //    BuildTree(transactions, -1, 0.5);
  void BuildTree(Transactions<T>* transactions,
                 int min_support_count,
                 double min_support_ratio);

  // Get the frequent itemsets using either min_support_count or
  // min_support_ratio. Only one min_support shall be properly set.
  // The other is set to any negative value.
  //
  // For example:
  // 1) Set by min_support_count = 10
  //    GetFrequentItemsets(10, -1, results);
  // 2) Set by min_support_ratio = 0.5
  //    GetFrequentItemsets(-1, 0.5, results);
  //
  // Note:
  //   The setting of min_support_count/ratio must be equal or greater
  //   than the same parameter in BuildTree.
  void GetFrequentItemsets(int min_support_count,
                           double min_support_ratio,
                           Itemsets<T>* results);

  // Version of the above method using a PatternWriter instance as output.
  void GetFrequentItemsets(int min_support_count,
                           double min_support_ratio,
                           PatternWriter<T>*
                           results);

  // Print the frequent-pattern tree stored internally for debug purpose.
  // The format of each node is [Term:Frequency:Index]
  // An example of frequent pattern tree is shown as follow.
  //
  // ____[:0:-1]____[mine:6:0]____[shoes:1:2]____[like:1:3]____[I:1:6]
  //           |             |____[tshirts:4:1]____[A:1:7]
  //           |             |
  //           |____[shoes:1:2]____[B:1:9]
  //           |              |____[pants:1:5]
  //           |____[irrelevant:1:4]____[totally:1:8]
  string TreeToString() const;

  // Print the Header Table for debug.
  string HeaderToString() const;

  // Print the dictionary for debug.
  string DictionaryToString() const;

  const ItemDictionary<T>& dictionary() const;

 private:

  // Internal function for recursively obtaining frequent pattern itemset
  void GetFrequentItemsetInternal(int frequency,
                                  Node* root,
                                  Header* header,
                                  int min_support_count,
                                  vector<int>* itemset,
                                  PatternWriter<T>* results);

  // Get the combination of itemsets on a single path
  void GetSinglePathItemset(const vector<int>& base_itemset,
                            int base_frequency,
                            int begin_item,
                            vector<int>* itemset,
                            PatternWriter<T>* results);

  // Given a frequent pattern tree, header table and an index item,
  // construct the conditional pattern base (itemsets and frequency)
  void ConditionalPatternBase(Node* tree,
                              Header* header,
                              int index,
                              vector<vector<int> >* base_itemsets,
                              vector<int>* base_frequencies);

  // Given frequent pattern base (itemsets and their frequency),
  // construct the corresponding conditional frequent pattern tree and header.
  void ConditionalPatternTree(vector<vector<int> >* base_itemsets,
                              vector<int>* base_frequencies,
                              int min_support_count,
                              Node* root,
                              Header* header);

  // Insert a transaction into a frequent pattern tree.
  // The transaction is an integer itemset specified by 'indices' with support
  // 'frequency'.
  void InsertToTree(const vector<int>& indices, int frequency);

  // Insert index itemset with a frequency into a frequent pattern tree
  void InsertToTreeInternal(const vector<int>& indices,
                            int frequency,
                            Node* root,
                            Header* header);

  // The internal function for printing the frequent pattern tree
  void TreeToStringInternal(const string& prefix,
                            Node* parent,
                            string* result) const;

  // Convert contents of 'pattern' from the internal indexed format into
  // the original items and append it to 'output'.
  void EmitPattern(const ItemsetFrequencyPair<int>& pattern,
                   PatternWriter<T>* output);

  // The root node of a frequent pattern tree.
  // Each node in the tree stores index rather than the original item.
  scoped_ptr<Node> root_;

  // The header table of a frequent pattern tree.
  scoped_ptr<Header> pattern_header_;

  // The dictionary index the mapping between items and indices.
  scoped_ptr<ItemDictionary<T> > dictionary_;

  // Minimal support for frequent pattern mining.
  // The valid range is in [0, num_transactions]
  // min_support_count = min_support_ratio * ceil(num_transactions)
  int min_support_count_;

  DISALLOW_EVIL_CONSTRUCTORS(Tree);
};

// ------------------------------------------------------------
//                      Definition
// ------------------------------------------------------------

// Constructor
template <typename T>
Tree<T>::Tree(T impossible_item)
  : root_(new Node(ItemDictionary<int>::IMPOSSIBLE_INDEX)),
    pattern_header_(new Header),
    dictionary_(new ItemDictionary<T>(impossible_item)) {
}

// Destructor
template <typename T>
Tree<T>::~Tree() {
}

// Insert a transaction into a frequent pattern tree.
// The transaction is an integer itemset specified by 'indices' with support
// 'frequency'.
template <typename T>
void Tree<T>::InsertToTree(const vector<int>& indices, int frequency) {
  InsertToTreeInternal(indices, frequency, root_.get(), pattern_header_.get());
}

// Insert index itemset with a frequency into a frequent pattern tree
template <typename T>
void Tree<T>::InsertToTreeInternal(const vector<int>& indices,
                                   int frequency,
                                   Node* root,
                                   Header* header) {
  Node* parent = root;
  for (int j = 0; j < indices.size(); ++j) {
    Node* node = NULL;
    for (int k = 0; k < parent->children().size(); ++k) {
      if (parent->children()[k]->index() == indices[j]) {
        node = parent->children()[k];
        break;
      }
    }
    if (node == NULL) {
      node = new Node(indices[j]);
      node->set_index(indices[j]);
      node->set_frequency(frequency);
      node->set_parent(parent);
      parent->add_child(node);
      header->add_node(node);
      header->add_frequency(indices[j], node->frequency());
    } else {
      node->set_frequency(node->frequency() + frequency);
      header->add_frequency(indices[j], frequency);
    }
    parent = node;
  }  // for each item in the transaction (j)
}

// Given a frequent pattern tree, header table and an index item,
// construct the conditional pattern base (itemsets and frequency)
// Note: a transaction with an empty itemset is allowed and necessary
template <typename T>
void Tree<T>::ConditionalPatternBase(
                            Node* tree,
                            Header* header,
                            int index,
                            vector<vector<int> >* base_itemsets,
                            vector<int>* base_frequencies) {
  Node* node = header->header_node(index);
  while (node != NULL) {
    int count = node->frequency();
    vector<int> transaction;
    // the transaction is empty if the node is one level down the root
    Node* path = node->parent();
    while (path != tree) {
      transaction.push_back(path->index());
      path = path->parent();
    }
    base_itemsets->push_back(transaction);
    base_frequencies->push_back(count);
    node = node->header_next();
  }
}

// Given frequent pattern base (itemsets and their frequency),
// construct the corresponding conditional frequent pattern tree and header.
template <typename T>
void Tree<T>::ConditionalPatternTree(
                            vector<vector<int> >* base_itemsets,
                            vector<int>* base_frequencies,
                            int min_support_count,
                            Node* root,
                            Header* header) {
  CHECK_EQ(base_itemsets->size(), base_frequencies->size());
  ItemDictionary<int> dictionary(ItemDictionary<int>::IMPOSSIBLE_INDEX);
  scoped_ptr<Transactions<int> > base_transactions(
    new TransactionsMemoryImpl<int>(base_itemsets, base_frequencies));
  dictionary.BuildDictionary(base_transactions.get(),
                             min_support_count,
                             -1);
  for (int i = 0; i < base_itemsets->size(); ++i) {
    vector<int> sorted_items;
    // Get Index from Item (also index)
    for (int j = 0; j < (*base_itemsets)[i].size(); ++j) {
      int index = dictionary.GetIndexForItem((*base_itemsets)[i][j]);
      if (index != dictionary.IMPOSSIBLE_INDEX) {
        sorted_items.push_back(index);
      }
    }
    sort(sorted_items.begin(), sorted_items.end());
    // Get Item (index) for Index
    for (int j = 0; j < sorted_items.size(); ++j) {
      sorted_items[j] = dictionary.GetItemForIndex(
                              sorted_items[j]);
    }
    InsertToTreeInternal(sorted_items,
                         (*base_frequencies)[i],
                         root,
                         header);
  }
}

// Get the combination of itemsets on a single path
template <typename T>
void Tree<T>::GetSinglePathItemset(
                            const vector<int>& base_itemset,
                            int base_frequency,
                            int begin_item,
                            vector<int>* itemset,
                            PatternWriter<T>* results) {
  ItemsetFrequencyPair<int> value;
  value.itemset = *itemset;
  value.frequency = base_frequency;
  EmitPattern(value, results);
  if (begin_item >= base_itemset.size()) {
    return;
  }
  for (int i = begin_item; i < base_itemset.size(); ++i) {
    itemset->push_back(base_itemset[i]);
    GetSinglePathItemset(base_itemset,
                         base_frequency,
                         i + 1,
                         itemset,
                         results);
    itemset->pop_back();
  }
}

// Internal function for constructing frequet itemsets.
// It enumerates all items in the tree and
//   1) build the conditional pattern base for a given item
//   2) build the conditional pattern tree from the pattern base.
//   3) recursively go to each sub tree and repeat step 1 & 2.
template <typename T>
void Tree<T>::GetFrequentItemsetInternal(
                            int frequency,
                            Node* root,
                            Header* header,
                            int min_support_count,
                            vector<int>* itemset,
                            PatternWriter<T>* results) {
  CHECK(root) << ": root can not be NULL";
  CHECK(header) << ": header can not be NULL";
  CHECK(itemset) << ": itemset can not be NULL";
  CHECK(results) << ": results can not be NULL";

  if (frequency > 0) {
    ItemsetFrequencyPair<int> value;
    value.itemset = *itemset;
    value.frequency = frequency;
    EmitPattern(value, results);
  }

  if (root->children().size() == 0) {
    return;
  }

  for (FrequencyTable::const_iterator i = header->frequency_table().begin();
       i != header->frequency_table().end();
       ++i) {
    // Skip the item if the frequency is less than min_support_count
    if (i->second < min_support_count) {
      continue;
    }
    // construct conditional pattern base
    vector<vector<int> >* base_itemsets =
      new vector<vector<int> >;
    vector<int>* base_frequencies =
      new vector<int>;
    LOG(INFO) << " the given item: "
            << dictionary_->GetItemTextForIndex(i->first);
    LOG(INFO) << " conditional pattern base";
    ConditionalPatternBase(root,
                           header,
                           i->first,
                           base_itemsets,
                           base_frequencies);
    if (base_itemsets->size() == 0) {
      LOG(WARNING) << ": the frequent pattern mining algorithm is problematic.";
      delete base_itemsets;
      delete base_frequencies;
    } else if (base_itemsets->size() == 1) {
      LOG(INFO) << ": frequent pattern base: single path or no path detected";
      itemset->push_back(i->first);
      GetSinglePathItemset((*base_itemsets)[0],
                           (*base_frequencies)[0],
                           0,
                           itemset,
                           results);
      itemset->pop_back();
      delete base_itemsets;
      delete base_frequencies;
    } else {
      // construct conditional pattern tree
      scoped_ptr<Node> conditional_root;
      conditional_root.reset(new Node(ItemDictionary<int>::IMPOSSIBLE_INDEX));
      scoped_ptr<Header> conditional_header;
      conditional_header.reset(new Header);
      LOG(INFO) << " conditional pattern tree";
      ConditionalPatternTree(base_itemsets,
                             base_frequencies,
                             min_support_count,
                             conditional_root.get(),
                             conditional_header.get());
      delete base_itemsets;
      delete base_frequencies;
      // DebugPrintTreeInternal("", conditional_root);
      itemset->push_back(i->first);
      // Recursively get the frequent pattern itemset
      GetFrequentItemsetInternal(i->second,
                                 conditional_root.get(),
                                 conditional_header.get(),
                                 min_support_count,
                                 itemset,
                                 results);
      itemset->pop_back();
    }  // if base_itemsets->size
  }  // for i
}

// Print the frequent-pattern tree stored internally, mainly for debug purpose.
// The format of each node is [Term:Frequency:Index]
// An example of frequent pattern tree is shown as follow.
//
// ____[:0:-1]____[mine:6:0]____[shoes:1:2]____[like:1:3]____[I:1:6]
//           |             |____[tshirts:4:1]____[A:1:7]
//           |             |
//           |____[shoes:1:2]____[B:1:9]
//           |             |____[pants:1:5]
//           |____[irrelevant:1:4]____[totally:1:8]
template <typename T>
string Tree<T>::TreeToString() const {
  LOG(INFO) << ": Print Frequent Pattern Tree";
  LOG(INFO) << ": Node format [Term:Frequency:Index]";
  string result;
  TreeToStringInternal("", root_.get(), &result);
  return result;
}

// The internal function for printing the frequent pattern tree
template <typename T>
void Tree<T>::TreeToStringInternal(const string& prefix,
                                   Node* parent,
                                   string* result) const {
  CHECK(parent) << ": parent can not be NULL";
  CHECK(result) << ": result can not be NULL";
  // output the node information
  string text = StringPrintf(
                  "____[%s:%d:%d]",
                  dictionary_->GetItemTextForIndex(parent->index()).c_str(),
                  parent->frequency(),
                  parent->index());
  *result += text;
  // build the padding prefix
  int padding_size = text.size();
  string padding_text;
  for (int i = 0; i < padding_size - 1; ++i) {
    padding_text += " ";
  }
  string new_prefix_0 = prefix + padding_text + "|";
  string new_prefix_1 = prefix + padding_text + " ";
  // go the children and show the branches
  if (parent->children().size() > 0) {
    if (parent->children().size() == 1) {
      // number of children == 1
      TreeToStringInternal(new_prefix_1, parent->children()[0], result);
    } else {
      // number of children >= 2
      TreeToStringInternal(new_prefix_0, parent->children()[0], result);
      // all the children except the first and last children.
      for (int i = 1; i < parent->children().size() - 1; ++i) {
          *result += new_prefix_0;
          TreeToStringInternal(new_prefix_0, parent->children()[i], result);
      }
      // the last child
      *result += new_prefix_0;
      TreeToStringInternal(new_prefix_1,
                           parent->children()[parent->children().size() -1],
                           result);
    }
  } else {
    *result += "\n";
  }
}

// Debug Print the Header Table for debug
template <typename T>
string Tree<T>::HeaderToString() const {
  string result;
  for (FrequencyTable::const_iterator i =
         pattern_header_->frequency_table().begin();
      i != pattern_header_->frequency_table().end();
      ++i) {
    result += StringPrintf("ITEM %s FREQUENCY %d\n",
                           dictionary_->GetItemTextForIndex(i->first).c_str(),
                           i->second);
  }
  return result;
}

// Print the dictionary for debug
template <typename T>
string Tree<T>::DictionaryToString() const {
  return dictionary_->ToString();
}

// Build the Frequent Pattern Tree using either min_support_count or
// min_support_ratio. Only one min_support shall be properly set.
// The other is set to any negative value. For example:
// 1) Set by min_support_count = 10
//    BuildTree(transactions, 10, -1);
// 2) Set by min_support_ratio = 0.5
//    BuildTree(transactions, -1, 0.5);
template <typename T>
void Tree<T>::BuildTree(Transactions<T>* transactions,
                        int min_support_count,
                        double min_support_ratio) {
  // There shall be only one valid min_support.
  CHECK_EQ(true, (min_support_count < 0 && min_support_ratio >= 0) ||
                 (min_support_count >= 0 && min_support_ratio < 0));

  // Build item dictionary:
  //   term-to-<index, frequency> mapping
  //   index-to-<term, frequency> mapping
  LOG(INFO) << ": Begin to build dictionary";
  dictionary_->BuildDictionary(transactions,
                               min_support_count,
                               min_support_ratio);
  min_support_count_ = min_support_count >= 0 ? min_support_count :
    GetCountFromRatio(dictionary_->num_transactions(), min_support_ratio);
  LOG_IF(WARNING, min_support_count_ > dictionary_->num_transactions())
    << "min_support_count shall be <= num_transactions. But it doesnt hurt.";
  LOG(INFO) << ": Finished building dictionary";

  // Run FP-tree algorithm to generate the frequent pattern tree
  LOG(INFO) << ": Begin to build FP-Tree";
  for (int i = 0; transactions->is_valid(i); ++i) {
    vector<T>& transaction = (*transactions)[i];
    if (transaction.size() == 0) {
      continue;
    }
    vector<int> indices;
    dictionary_->ConvertTransactionToSortedIndices(transaction, &indices);
    InsertToTree(indices, (*transactions)(i));
  }  // for all transactions (i)
  LOG(INFO) << ": Finished building FP-Tree";
}

template <typename T>
void Tree<T>::EmitPattern(const ItemsetFrequencyPair<int>& pattern,
                          PatternWriter<T>* output) {
  // Order item indices so that their order keeps the same as in the original
  // transactions statistically, if needed.
  vector<int> indexed_itemset = pattern.itemset;
  if (FLAGS_keep_item_order) {
    dictionary_->SortIndicesByPartialOrder(&indexed_itemset);
  }
  ItemsetFrequencyPair<T> value;
  value.frequency = pattern.frequency;
  value.is_partial = false;
  value.is_phrase = false;
  value.itemset.reserve(indexed_itemset.size());
  for (int j = 0; j < indexed_itemset.size(); ++j) {
    value.itemset.push_back(
        dictionary_->GetItemForIndex(indexed_itemset[j]));
  }
  output->WritePattern(value);
}

// Get the frequent itemsets using either min_support_count or
// min_support_ratio. Only one min_support shall be properly set.
// The other is set to any negative value.
//
// For example:
// 1) Set by min_support_count = 10
//    GetFrequentItemsets(10, -1, results);
// 2) Set by min_support_ratio = 0.5
//    GetFrequentItemsets(-1, 0.5, results);
//
// Note:
//   The setting of min_support_count/ratio must be equal or greater
//   than the same parameter in BuildTree.
template <typename T>
void Tree<T>::GetFrequentItemsets(int min_support_count,
                                  double min_support_ratio,
                                  Itemsets<T>* results) {
  scoped_ptr<PatternWriterMemoryImpl<T> > writer(
      new PatternWriterMemoryImpl<T>(results));
  GetFrequentItemsets(min_support_count, min_support_ratio, writer.get());

  LOG(INFO) << ": Mined a total of " << results->itemsets().size()
            << " patterns";

  results->SortByItemsetSizeFrequency();
  // Set the number of transactions.
  results->set_num_transactions(dictionary_->num_transactions());
}

template <typename T>
void Tree<T>::GetFrequentItemsets(int min_support_count,
                                  double min_support_ratio,
                                  PatternWriter<T>* results) {
  // Either min_support_count or min_support_ratio is set to a proper value.
  CHECK_EQ(true, (min_support_count < 0 && min_support_ratio >= 0) ||
                 (min_support_count >= 0 && min_support_ratio < 0));

  // In BuildTree(), we have a requirement for min_support_count
  // The count in this function must be >= the count in BuildTree()
  min_support_count = min_support_count >= 0 ? min_support_count :
    GetCountFromRatio(dictionary_->num_transactions(), min_support_ratio);
  LOG_IF(WARNING, min_support_count > dictionary_->num_transactions())
    << "min_support_count shall be <= num_transactions thought it doesnt hurt.";

  CHECK_GE(min_support_count, min_support_count_)
    << ": min_support_count must be >= that for building tree and dictionary";

  CHECK(results) << ": results can not be NULL";

  vector<int> itemset;
  GetFrequentItemsetInternal(0,
                             root_.get(),
                             pattern_header_.get(),
                             min_support_count,
                             &itemset,
                             results);
}

// Retrieve the ItemDictionary.
template <typename T>
const ItemDictionary<T>& Tree<T>::dictionary() const {
  return *dictionary_;
}

}  // namespace frequent_pattern_mining

#endif  // FP_GROWTH_TREE_H_
